#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
PyTorch 环境验证脚本
检查 PyTorch 安装和 GPU 支持情况
"""

import torch

def check_pytorch_installation():
    """检查 PyTorch 安装状态"""
    print("=" * 50)
    print("PyTorch 环境验证")
    print("=" * 50)

    # 1. 检查 PyTorch 版本
    try:
        print(f"PyTorch 版本: {torch.__version__}")
    except ImportError:
        print("❌ PyTorch 未安装")
        return False

    # 2. 检查 CUDA 可用性
    cuda_available = torch.cuda.is_available()
    print(f"CUDA 可用: {cuda_available}")

    if cuda_available:
        # 3. 检查 GPU 数量和名称
        gpu_count = torch.cuda.device_count()
        print(f"GPU 数量: {gpu_count}")

        for i in range(gpu_count):
            gpu_name = torch.cuda.get_device_name(i)
            print(f"GPU {i}: {gpu_name}")

        # 4. 检查 CUDA 版本
        cuda_version = torch.version.cuda
        print(f"CUDA 版本: {cuda_version}")

        # 5. 检查 GPU 内存
        for i in range(gpu_count):
            total_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
            print(f"GPU {i} 总内存: {total_memory:.1f} GB")

        print("✅ PyTorch GPU 支持正常！")
        return True
    else:
        print("❌ CUDA 不可用，请检查 GPU 驱动和 PyTorch 版本")
        return False

def test_basic_operations():
    """测试基本 PyTorch 操作"""
    print("\n" + "=" * 50)
    print("基本操作测试")
    print("=" * 50)

    try:
        # CPU 测试
        x = torch.randn(10, 10)
        y = torch.randn(10, 10)
        z = torch.matmul(x, y)
        print(f"CPU 矩阵乘法成功，结果形状: {z.shape}")

        # GPU 测试（如果可用）
        if torch.cuda.is_available():
            x_gpu = x.cuda()
            y_gpu = y.cuda()
            z_gpu = torch.matmul(x_gpu, y_gpu)
            print(f"GPU 矩阵乘法成功，结果形状: {z_gpu.shape}")

        print("✅ 基本操作测试通过！")
        return True
    except Exception as e:
        print(f"❌ 基本操作测试失败: {e}")
        return False

if __name__ == "__main__":
    print("开始验证 PyTorch 环境...")

    # 检查安装
    install_ok = check_pytorch_installation()

    # 测试基本操作
    if install_ok:
        test_ok = test_basic_operations()

        if test_ok:
            print("\n🎉 PyTorch 环境验证完成，一切正常！")
        else:
            print("\n⚠️  PyTorch 环境存在问题，需要进一步检查")
    else:
        print("\n❌ PyTorch 环境验证失败，请重新安装")